#!/bin/bash

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'

# 日志函数
log_info() {
    echo -e "${GREEN}[INFO] $1${NC}"
}

log_warn() {
    echo -e "${YELLOW}[WARN] $1${NC}"
}

log_error() {
    echo -e "${RED}[ERROR] $1${NC}"
}

# 检查必要命令
check_command() {
    if ! command -v $1 &> /dev/null; then
        log_error "未找到命令: $1"
        exit 1
    fi
}

# 检查 Trino 服务状态
check_trino_status() {
    if ! systemctl is-active --quiet trino; then
        log_error "Trino 服务未运行"
        exit 1
    fi
}

# 添加 Worker 节点
add_worker() {
    local node_name=$1
    local node_ip=$2
    local coordinator_ip=$3
    local coordinator_port=$4

    log_info "开始添加 Worker 节点: $node_name ($node_ip)"

    # 检查节点连通性
    if ! ssh -o BatchMode=yes -o ConnectTimeout=5 $node_ip echo &> /dev/null; then
        log_error "无法连接到节点: $node_ip"
        exit 1
    fi

    # 在远程节点安装 Trino
    log_info "在节点 $node_name 上安装 Trino..."
    ssh $node_ip "sudo apt-get update && sudo apt-get install -y trino"

    # 配置 Trino
    log_info "配置 Trino..."
    ssh $node_ip "sudo mkdir -p /etc/trino"
    ssh $node_ip "sudo tee /etc/trino/config.properties > /dev/null << EOF
coordinator=false
http-server.http.port=8080
discovery.uri=http://$coordinator_ip:$coordinator_port
EOF"

    # 启动服务
    log_info "启动 Trino 服务..."
    ssh $node_ip "sudo systemctl enable trino && sudo systemctl start trino"

    # 验证节点是否成功加入
    sleep 10
    if curl -s "http://$coordinator_ip:$coordinator_port/v1/cluster" | grep -q "$node_name"; then
        log_info "Worker 节点 $node_name 已成功添加到集群"
    else
        log_error "Worker 节点添加失败"
        exit 1
    fi
}

# 移除 Worker 节点
remove_worker() {
    local node_name=$1
    local node_ip=$2
    local coordinator_ip=$3
    local coordinator_port=$4

    log_info "开始移除 Worker 节点: $node_name ($node_ip)"

    # 检查节点是否在集群中
    if ! curl -s "http://$coordinator_ip:$coordinator_port/v1/cluster" | grep -q "$node_name"; then
        log_error "节点 $node_name 不在集群中"
        exit 1
    fi

    # 停止服务
    log_info "停止 Trino 服务..."
    ssh $node_ip "sudo systemctl stop trino"

    # 清理数据
    log_info "清理节点数据..."
    ssh $node_ip "sudo rm -rf /var/lib/trino/data"

    # 验证节点是否已移除
    sleep 10
    if ! curl -s "http://$coordinator_ip:$coordinator_port/v1/cluster" | grep -q "$node_name"; then
        log_info "Worker 节点 $node_name 已成功从集群移除"
    else
        log_error "Worker 节点移除失败"
        exit 1
    fi
}

# 添加磁盘
add_disk() {
    local node_name=$1
    local node_ip=$2
    local disk_path=$3

    log_info "开始为节点 $node_name 添加磁盘: $disk_path"

    # 检查节点是否在集群中
    if ! curl -s "http://$coordinator_ip:$coordinator_port/v1/cluster" | grep -q "$node_name"; then
        log_error "节点 $node_name 不在集群中"
        exit 1
    fi

    # 创建数据目录
    log_info "创建数据目录..."
    ssh $node_ip "sudo mkdir -p $disk_path/trino/data && sudo chown -R trino:trino $disk_path/trino"

    # 更新配置
    log_info "更新 Trino 配置..."
    ssh $node_ip "sudo tee -a /etc/trino/config.properties > /dev/null << EOF
data-directory=$disk_path/trino/data
EOF"

    # 重启服务
    log_info "重启 Trino 服务..."
    ssh $node_ip "sudo systemctl restart trino"

    # 验证磁盘是否成功添加
    sleep 10
    if ssh $node_ip "df -h $disk_path" &> /dev/null; then
        log_info "磁盘已成功添加到节点 $node_name"
    else
        log_error "磁盘添加失败"
        exit 1
    fi
}

# 处理故障磁盘
handle_failed_disk() {
    local node_name=$1
    local node_ip=$2
    local disk_path=$3

    log_info "开始处理节点 $node_name 的故障磁盘: $disk_path"

    # 检查节点是否在集群中
    if ! curl -s "http://$coordinator_ip:$coordinator_port/v1/cluster" | grep -q "$node_name"; then
        log_error "节点 $node_name 不在集群中"
        exit 1
    fi

    # 从配置中移除磁盘
    log_info "从配置中移除磁盘..."
    ssh $node_ip "sudo sed -i '/data-directory=$disk_path/d' /etc/trino/config.properties"

    # 重启服务
    log_info "重启 Trino 服务..."
    ssh $node_ip "sudo systemctl restart trino"

    # 卸载磁盘
    log_info "卸载故障磁盘..."
    ssh $node_ip "sudo umount $disk_path"

    # 标记磁盘为故障
    log_info "标记磁盘为故障状态..."
    ssh $node_ip "sudo echo 'FAILED' > $disk_path/disk_status"

    log_info "故障磁盘处理完成"
}

# 主函数
main() {
    # 检查必要命令
    check_command ssh
    check_command curl
    check_command systemctl

    # 检查 Trino 服务状态
    check_trino_status

    # 解析命令行参数
    case "$1" in
        "add-worker")
            if [ $# -ne 5 ]; then
                log_error "用法: $0 add-worker <node_name> <node_ip> <coordinator_ip> <coordinator_port>"
                exit 1
            fi
            add_worker "$2" "$3" "$4" "$5"
            ;;
        "remove-worker")
            if [ $# -ne 5 ]; then
                log_error "用法: $0 remove-worker <node_name> <node_ip> <coordinator_ip> <coordinator_port>"
                exit 1
            fi
            remove_worker "$2" "$3" "$4" "$5"
            ;;
        "add-disk")
            if [ $# -ne 4 ]; then
                log_error "用法: $0 add-disk <node_name> <node_ip> <disk_path>"
                exit 1
            fi
            add_disk "$2" "$3" "$4"
            ;;
        "handle-failed-disk")
            if [ $# -ne 4 ]; then
                log_error "用法: $0 handle-failed-disk <node_name> <node_ip> <disk_path>"
                exit 1
            fi
            handle_failed_disk "$2" "$3" "$4"
            ;;
        *)
            log_error "用法: $0 {add-worker|remove-worker|add-disk|handle-failed-disk}"
            exit 1
            ;;
    esac
}

# 执行主函数
main "$@" 